from flask import Blueprint, jsonify, request
from models.document import Knowledge, Document, DocumentSegment, db
from services.ollama_service import OllamaService
#from services.vector_service import VectorService
from services.chroma_service import ChromaService
import uuid

knowledge_bp = Blueprint('knowledge', __name__)
ollama_service = OllamaService()
chroma_service = ChromaService()

@knowledge_bp.route('/models', methods=['GET'])
def get_models():
    """获取可用的嵌入模型列表"""
    try:
        models = ollama_service.list_models()
        return jsonify(models)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@knowledge_bp.route('/', methods=['GET'])
def list_knowledges():
    """获取知识库列表"""
    try:
        knowledges = Knowledge.query.all()
        result = []
        for k in knowledges:
            doc_count = Document.query.filter_by(knowledge_id=k.id).count()
            # segment_count = db.session.query(DocumentSegment).\
            #     join(Document).\
            #     filter(Document.knowledge_id == k.id).count()
            segment_count = chroma_service.get_count(k.collection_name)
            
            result.append({
                "id": k.id,
                "name": k.name,
                "embedding_model": k.embedding_model,
                "doc_count": doc_count,
                "segment_count": segment_count,
                "collection_name": k.collection_name
            })
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@knowledge_bp.route('/', methods=['POST'])
def create_knowledge():
    """创建新知识库"""
    try:
        data = request.json
        collection_name = data.get('collection_name') or f"k_{uuid.uuid4().hex[:8]}"
        
        knowledge = Knowledge(
            name=data['name'],
            embedding_model=data['embedding_model'],
            collection_name=collection_name
        )
        
        db.session.add(knowledge)
        db.session.commit()
        
        # 创建向量集合
        chroma_service.create_collection(collection_name)
        
        return jsonify({"id": knowledge.id})
    except Exception as e:
        db.session.rollback()
        return jsonify({"error": str(e)}), 500

@knowledge_bp.route('/<int:knowledge_id>', methods=['PATCH'])
def update_knowledge(knowledge_id):
    """更新知识库信息"""
    try:
        knowledge = Knowledge.query.get_or_404(knowledge_id)
        data = request.json
        
        if 'name' in data:
            knowledge.name = data['name']
            
        db.session.commit()
        
        return jsonify({
            "id": knowledge.id,
            "name": knowledge.name,
            "embedding_model": knowledge.embedding_model,
            "collection_name": knowledge.collection_name
        })
        
    except Exception as e:
        db.session.rollback()
        return jsonify({"error": str(e)}), 500

@knowledge_bp.route('/<int:knowledge_id>/search', methods=['POST'])
def search_knowledge(knowledge_id):
    """检索知识库内容"""
    try:
        query = request.json['query']
        knowledge = Knowledge.query.get_or_404(knowledge_id)
        
        # 首先获取查询文本的向量表示
        # query_embedding = ollama_service.get_embeddings(query)
        
        # 获取向量检索结果
        results = chroma_service.list_records(
            collection_name=knowledge.collection_name,
            page=1,
            page_size=10,
            search_type='hybrid',
            record_id=None,
            query_text=query,
            metadata_filter=None,
            document_filter=None,
            n_results_requested=10
        )

        if not results.get('records'):  # 检查是否存在有效记录集
            return []

        # 查询对应的文档信息并处理结果
        response = []
        for record in results['records']:
            try:
                if 'metadata' in record and 'document_id' in record['metadata']:
                    document_id = int(record['metadata']['document_id'])
                    document = Document.query.get(document_id)
                    # 计算相似度得分 (将距离转换为0-100的得分)
                    similarity_score = max(0, min(100, (1 - record.get('distance', 0)) * 100))
                    response.append({
                        "content": record['document'],
                        "score": round(similarity_score, 2),  # 保留两位小数
                        "filename": document.filename if document else "未知文档",
                        "metadata": record['metadata']
                    })
            except (KeyError, ValueError, TypeError) as e:
                continue
            
        return jsonify(response)
    except Exception as e:
        return jsonify({"error": str(e)}), 500